import torch
import torch.utils.data as data
#import torchvision.transforms as transforms
import os
import nltk
from PIL import Image
import numpy as np
import json as jsonmod


class PrecompDataset(data.Dataset):
    """
    Load precomputed captions and image features  # 加载预先计算的图像和标题特征
    Possible options: f30k_precomp, coco_precomp
    """

    def __init__(self, data_path, data_split, vocab):
        self.vocab = vocab  # 这里的vocab是一个类
        loc = data_path + '/'  # 这里是数据的存储路径

        # Captions
        self.captions = []
        with open(loc+'%s_caps.txt' % data_split, 'rb') as f:  # path/train_caps.txt
            # 从txt里按行读取标题信息，存储到列表中
            for line in f:
                self.captions.append(line.strip())

        # Image features
        self.images = np.load(loc+'%s_ims.npy' % data_split)  # path/train_ims.npy
        self.img_boxes = np.load(loc+'%s_boxes.npy' % data_split)
        self.length = len(self.captions)  # 为标题长度赋值
        # rkiros data has redundancy in images, we divide by 5, 10crop doesn't
        if self.images.shape[0] != self.length:  # 如果图像的行数不等于标题的行数
            self.im_div = 5
        else:
            self.im_div = 1
        # the development set for coco is large and so validation would be slow
        if data_split == 'dev':  # 猜测这里是防止验证时速度慢，所以只取5000个样例
            self.length = 5000
            if not torch.cuda.is_available():
                self.length = 25
                print('LPY: NO CUDA DEVICE!!! DEV ONLY 25, PC DEBUG NUM.')

    # 如果在类中定义了__getitem__()方法，那么他的实例对象（假设为P）就可以这样P[key]取值。
    # 当实例对象做P[key]运算时，就会调用类中的__getitem__()方法。
    # https://blog.csdn.net/chituozha5528/article/details/78354833
    def __getitem__(self, index):
        # handle the image redundancy
        img_id = index//self.im_div
        image = torch.Tensor(self.images[img_id])
        imgbox = torch.Tensor(self.img_boxes[img_id])
        caption = self.captions[index]
        vocab = self.vocab

        # Convert caption (string) to word ids.
        #tokens = nltk.tokenize.word_tokenize(
         #   str(caption).lower().decode('utf-8'))
        tokens = nltk.tokenize.word_tokenize(str(caption).lower())
        caption = []
        caption.append(vocab('<start>'))
        caption.extend([vocab(token) for token in tokens])
        caption.append(vocab('<end>'))
        target = torch.Tensor(caption)
        return image, target, index, img_id,imgbox

    def __len__(self):
        return self.length


def collate_fn(data):
    """Build mini-batch tensors from a list of (image, caption) tuples.
    Args:
        data: list of (image, caption) tuple.
            - image: torch tensor of shape (3, 256, 256).
            - caption: torch tensor of shape (?); variable length.

    Returns:
        images: torch tensor of shape (batch_size, 3, 256, 256).
        targets: torch tensor of shape (batch_size, padded_length).
        lengths: list; valid length for each padded caption.
    """
    # Sort a data list by caption length
    data.sort(key=lambda x: len(x[1]), reverse=True)
    images, captions, ids, img_ids,img_boxes = zip(*data)

    # Merge images (convert tuple of 3D tensor to 4D tensor)
    images = torch.stack(images, 0)
    img_boxes = torch.stack(img_boxes,0)

    # Merget captions (convert tuple of 1D tensor to 2D tensor)
    lengths = [len(cap) for cap in captions]
    targets = torch.zeros(len(captions), max(lengths)).long()
    for i, cap in enumerate(captions):
        end = lengths[i]
        targets[i, :end] = cap[:end]

    return images, targets, lengths, ids, img_boxes


def get_precomp_loader(data_path, data_split, vocab, opt, batch_size=100,
                       shuffle=True, num_workers=2):
    """Returns torch.utils.data.DataLoader for custom coco dataset."""
    dset = PrecompDataset(data_path, data_split, vocab)

    # torch.utils.data.DataLoader 数据加载器，结合了数据集和取样器，并且可以提供多个线程处理数据集。
    # 在训练模型时使用到此函数，用来把训练数据分成多个小组，此函数每次抛出一组数据。
    # 直至把所有的数据都抛出。就是做一个数据的初始化。
    # https://www.cnblogs.com/demo-deng/p/10623334.html
    # 第一个参数dataset，应该处理成什么格式？待确认
    # 这里的dset是一个类，那么对这个类有什么要求呢？
    data_loader = torch.utils.data.DataLoader(dataset=dset,
                                              batch_size=batch_size,
                                              shuffle=shuffle,
                                              pin_memory=True,
                                              collate_fn=collate_fn)
    return data_loader

# 加载了数据，并且把数据按照batch_size分成了待训练的小组
def get_loaders(data_name, vocab, batch_size, workers, opt):
    dpath = os.path.join(opt.data_path, data_name)   # 拼接了数据的路径
    train_loader = get_precomp_loader(dpath, 'train', vocab, opt,
                                      batch_size, True, workers)
    if opt.dev_data=='test':
        val_loader = get_precomp_loader(dpath, 'test', vocab, opt,
                                        batch_size, False, workers)
    elif opt.dev_data=='dev':
        val_loader = get_precomp_loader(dpath, 'dev', vocab, opt,
                                    batch_size, False, workers)
    return train_loader, val_loader


def get_test_loader(split_name, data_name, vocab, batch_size,
                    workers, opt):
    dpath = os.path.join(opt.data_path, data_name)
    test_loader = get_precomp_loader(dpath, split_name, vocab, opt,
                                     batch_size, False, workers)
    return test_loader
